'''
性能瓶颈主要在探测文字区域
'''
import time

from cnstd import CnStd
from cnocr import CnOcr

from submodules.common.lib.log_ import logger
logger =get_logger()


def get_bounds_by_text(file_path, target_text=None, context="cpu"):
    res = []

    std = CnStd(context=context)
    cn_ocr = CnOcr(context=context)

    box_infos = std.detect(file_path)
    for box_info in box_infos['detected_texts']:
        x = box_info["box"][0]
        y = box_info["box"][1]
        w = box_info["box"][2]
        h = box_info["box"][3]
        x1 = x - w / 2
        y1 = y - h / 2
        x2 = x + w / 2
        y2 = y + h / 2
        cropped_img = box_info['cropped_img']
        ocr_res = cn_ocr.ocr_for_single_line(cropped_img)
        ocr_text = "".join(ocr_res[0])
        ocr_prob = ocr_res[1]
        if target_text == None:
            res.append([(x1, y1, x2, y2), ocr_text, ocr_prob])
        else:
            if ocr_text == target_text:
                res.append([(x1, y1, x2, y2), ocr_text, ocr_prob])
    return res


def ocr(file_path, context="cpu"):
    t1 = time.time()
    res = get_bounds_by_text(file_path=file_path, target_text=None, context=context)
    t2 = time.time()
    logger.info(f"ocr_time:{t2 - t1}")
    return res


if __name__ == "__main__":
    # res = cnocr_warpped("test_data/test.png")
    # print(res)

    def get_bounds_by_text_test():
        text = "同意"
        img_path = "test_data/test.png"
        res = get_bounds_by_text(img_path, text)
        print(res)
        from submodules.common.lib.image_ import img_cut
        img_cut(img_path, res[0][0])


    def test_ocr():
        img_path = "test_data/test.png"
        res = ocr(img_path)
        print(res)


    test_ocr()
